from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.prebuilt import create_react_agent
from langgraph.store.postgres import PostgresStore

from agent.my_llm import llm
from agent.tools.tools_demo6 import runnable_tool
from agent.tools.tools_demo7 import MySearchTool


# 创建一个网络搜索的工具
search_tool = MySearchTool()


# 短期记忆，保存到内存中
# checkpointer = InMemorySaver()

# langgraph dev 平台部署会报错
# ValueError: Heads up! Your graph 'agent' from './src/agent/my_agent2.py' includes a custom checkpointer (type
# <class 'langgraph.checkpoint.memory.InMemorySaver'>). With LangGraph API, persistence is handled automatically by
# the platform, so providing a custom checkpointer (type <class 'langgraph.checkpoint.memory.InMemorySaver'>) here
# isn't necessary and will be ignored when deployed.

# 在生产环境中，使用由数据库支持的检查点
# 必须安装：pip install -U "psycopg[binary,pool]" langgraph-checkpoint-postgres
# DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"
DB_URI = "postgresql://root:root123456..@localhost:5432/langgraph_db?sslmode=disable"
with (
    PostgresStore.from_conn_string(DB_URI) as store,
    PostgresSaver.from_conn_string(DB_URI) as checkpointer
):

    # 首次使用时调用，初始化数据库
    checkpointer.setup()
    store.setup()

    agent = create_react_agent(
        llm,
        tools=[runnable_tool, search_tool],
        prompt="你是一个智能助手，尽可能的调用工具回答用户的问题",
        checkpointer=checkpointer,
        store=store
    )

    config = {
        "configurable": {
            "thread_id": "111"
        }
    }

    # 从短期存储中，获取返回所有当前会话的上下文
    rest0 = list(agent.get_state(config))
    print(rest0)
    # 从长期存储中，获取返回所有当前会话的上下文
    rest1 = list(agent.get_state_history(config))
    print(rest1)


    # resp1 = agent.invoke(
    #     {"messages": [{"role": "user", "content": "今天上海的天气怎么样？"}]},
    #     config
    # )
    # print(resp1['messages'][-1].content)
    #
    # print('\n')
    #
    # resp2 = agent.invoke(
    #     {"messages": [{"role": "user", "content": "那北京呢？"}]},
    #     config
    # )
    # print(resp2['messages'][-1].content)




